from pymilvus import model
from pymilvus import MilvusClient



##第1行：创建一个默认的嵌入函数（embedding function）实例。嵌入函数是一种将文本转换为数字向量的工具，在AI和机器学习中经常用到。
##这行代码创建了一个模型对象，它可以将文字转换成计算机能理解的数字形式。
embedding_fn = model.DefaultEmbeddingFunction()

##第2-6行：定义一个文档列表（docs），里面包含了3个关于人工智能和图灵的英文句子。每个句子都是一个字符串元素，存储在这个列表中。
docs = [
    "Artificial intelligence was founded as an academic discipline in 1956.",
    "Alan Turing was the first person to conduct substantial research in AI.",
    "Born in Maida Vale, London, Turing was raised in southern England.",
]

##第7行：使用前面创建的嵌入函数将文档列表中的文本转换为向量。encode_documents方法会处理整个文档列表，将每条文本转换为对应的数字向量表示。
vectors = embedding_fn.encode_documents(docs)

##第8行：打印输出向量的维度信息。这里显示嵌入函数的维度是768，第一个向量的形状是(768,)，意味着每个文本都被转换成了一个包含768个数字的向量。
print("Dim:", embedding_fn.dim, vectors[0].shape)  # Dim: 768 (768,)
# print("Vector:", vectors)

##第9-12行：创建一个新的数据列表，其中包含每个文档的详细信息：

##id：文档的索引号（0, 1, 2）
##vector：对应的向量表示
##text：原始文本内容
##subject：主题分类，这里都标记为"history"
##这是一个列表推导式，它遍历所有向量，为每个文档创建一个包含上述字段的字典
data = [
    {"id": i, "vector": vectors[i], "text": docs[i], "subject": "history"}
    for i in range(len(vectors))
]

# print("Data:", data)

##第14-15行：打印数据集的基本信息：

##第一行显示数据集中有多少个实体（文档数量），以及每个实体包含哪些字段
##第二行显示向量的维度（即向量中包含多少个数值）
print("Data has", len(data), "entities, each with fields: ", data[0].keys())
print("Vector dim:", len(data[0]["vector"]))


client = MilvusClient("./milvus_demo.db")
if client.has_collection(collection_name="demo_collection"):
    client.drop_collection(collection_name="demo_collection")
client.create_collection(
    collection_name="demo_collection",
    dimension=768,  # The vectors we will use in this demo has 768 dimensions
)





res = client.insert(collection_name="demo_collection", data=data)

print(res)


query_vectors = embedding_fn.encode_queries(["Who is Alan Turing?"])

res = client.search(
    collection_name="demo_collection",  # target collection
    data=query_vectors,  # query vectors
    limit=2,  # number of returned entities
    output_fields=["text", "subject"],  # specifies fields to be returned
)

print(res)
